// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2012 Gael Guennebaud <gael.guennebaud@inria.fr>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_SPARSELU_GEMM_KERNEL_H
#define EIGEN_SPARSELU_GEMM_KERNEL_H

namespace Eigen {

namespace internal {

/** \internal
 * A general matrix-matrix product kernel optimized for the SparseLU factorization.
 *  - A, B, and C must be column major
 *  - lda and ldc must be multiples of the respective packet size
 *  - C must have the same alignment as A
 */
template<typename Scalar>
EIGEN_DONT_INLINE void
sparselu_gemm(Index m, Index n, Index d, const Scalar* A, Index lda, const Scalar* B, Index ldb, Scalar* C, Index ldc)
{
	using namespace Eigen::internal;

	typedef typename packet_traits<Scalar>::type Packet;
	enum
	{
		NumberOfRegisters = EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS,
		PacketSize = packet_traits<Scalar>::size,
		PM = 8,								  // peeling in M
		RN = 2,								  // register blocking
		RK = NumberOfRegisters >= 16 ? 4 : 2, // register blocking
		BM = 4096 / sizeof(Scalar),			  // number of rows of A-C per chunk
		SM = PM * PacketSize				  // step along M
	};
	Index d_end = (d / RK) * RK; // number of columns of A (rows of B) suitable for full register blocking
	Index n_end = (n / RN) * RN; // number of columns of B-C suitable for processing RN columns at once
	Index i0 = internal::first_default_aligned(A, m);

	eigen_internal_assert(((lda % PacketSize) == 0) && ((ldc % PacketSize) == 0) &&
						  (i0 == internal::first_default_aligned(C, m)));

	// handle the non aligned rows of A and C without any optimization:
	for (Index i = 0; i < i0; ++i) {
		for (Index j = 0; j < n; ++j) {
			Scalar c = C[i + j * ldc];
			for (Index k = 0; k < d; ++k)
				c += B[k + j * ldb] * A[i + k * lda];
			C[i + j * ldc] = c;
		}
	}
	// process the remaining rows per chunk of BM rows
	for (Index ib = i0; ib < m; ib += BM) {
		Index actual_b = std::min<Index>(BM, m - ib);				// actual number of rows
		Index actual_b_end1 = (actual_b / SM) * SM;					// actual number of rows suitable for peeling
		Index actual_b_end2 = (actual_b / PacketSize) * PacketSize; // actual number of rows suitable for vectorization

		// Let's process two columns of B-C at once
		for (Index j = 0; j < n_end; j += RN) {
			const Scalar* Bc0 = B + (j + 0) * ldb;
			const Scalar* Bc1 = B + (j + 1) * ldb;

			for (Index k = 0; k < d_end; k += RK) {

				// load and expand a RN x RK block of B
				Packet b00, b10, b20, b30, b01, b11, b21, b31;
				{
					b00 = pset1<Packet>(Bc0[0]);
				}
				{
					b10 = pset1<Packet>(Bc0[1]);
				}
				if (RK == 4) {
					b20 = pset1<Packet>(Bc0[2]);
				}
				if (RK == 4) {
					b30 = pset1<Packet>(Bc0[3]);
				}
				{
					b01 = pset1<Packet>(Bc1[0]);
				}
				{
					b11 = pset1<Packet>(Bc1[1]);
				}
				if (RK == 4) {
					b21 = pset1<Packet>(Bc1[2]);
				}
				if (RK == 4) {
					b31 = pset1<Packet>(Bc1[3]);
				}

				Packet a0, a1, a2, a3, c0, c1, t0, t1;

				const Scalar* A0 = A + ib + (k + 0) * lda;
				const Scalar* A1 = A + ib + (k + 1) * lda;
				const Scalar* A2 = A + ib + (k + 2) * lda;
				const Scalar* A3 = A + ib + (k + 3) * lda;

				Scalar* C0 = C + ib + (j + 0) * ldc;
				Scalar* C1 = C + ib + (j + 1) * ldc;

				a0 = pload<Packet>(A0);
				a1 = pload<Packet>(A1);
				if (RK == 4) {
					a2 = pload<Packet>(A2);
					a3 = pload<Packet>(A3);
				} else {
					// workaround "may be used uninitialized in this function" warning
					a2 = a3 = a0;
				}

#define KMADD(c, a, b, tmp)                                                                                            \
	{                                                                                                                  \
		tmp = b;                                                                                                       \
		tmp = pmul(a, tmp);                                                                                            \
		c = padd(c, tmp);                                                                                              \
	}
#define WORK(I)                                                                                                        \
	c0 = pload<Packet>(C0 + i + (I)*PacketSize);                                                                       \
	c1 = pload<Packet>(C1 + i + (I)*PacketSize);                                                                       \
	KMADD(c0, a0, b00, t0)                                                                                             \
	KMADD(c1, a0, b01, t1)                                                                                             \
	a0 = pload<Packet>(A0 + i + (I + 1) * PacketSize);                                                                 \
	KMADD(c0, a1, b10, t0)                                                                                             \
	KMADD(c1, a1, b11, t1)                                                                                             \
	a1 = pload<Packet>(A1 + i + (I + 1) * PacketSize);                                                                 \
	if (RK == 4) {                                                                                                     \
		KMADD(c0, a2, b20, t0)                                                                                         \
	}                                                                                                                  \
	if (RK == 4) {                                                                                                     \
		KMADD(c1, a2, b21, t1)                                                                                         \
	}                                                                                                                  \
	if (RK == 4) {                                                                                                     \
		a2 = pload<Packet>(A2 + i + (I + 1) * PacketSize);                                                             \
	}                                                                                                                  \
	if (RK == 4) {                                                                                                     \
		KMADD(c0, a3, b30, t0)                                                                                         \
	}                                                                                                                  \
	if (RK == 4) {                                                                                                     \
		KMADD(c1, a3, b31, t1)                                                                                         \
	}                                                                                                                  \
	if (RK == 4) {                                                                                                     \
		a3 = pload<Packet>(A3 + i + (I + 1) * PacketSize);                                                             \
	}                                                                                                                  \
	pstore(C0 + i + (I)*PacketSize, c0);                                                                               \
	pstore(C1 + i + (I)*PacketSize, c1)

				// process rows of A' - C' with aggressive vectorization and peeling
				for (Index i = 0; i < actual_b_end1; i += PacketSize * 8) {
					EIGEN_ASM_COMMENT("SPARSELU_GEMML_KERNEL1");
					prefetch((A0 + i + (5) * PacketSize));
					prefetch((A1 + i + (5) * PacketSize));
					if (RK == 4)
						prefetch((A2 + i + (5) * PacketSize));
					if (RK == 4)
						prefetch((A3 + i + (5) * PacketSize));

					WORK(0);
					WORK(1);
					WORK(2);
					WORK(3);
					WORK(4);
					WORK(5);
					WORK(6);
					WORK(7);
				}
				// process the remaining rows with vectorization only
				for (Index i = actual_b_end1; i < actual_b_end2; i += PacketSize) {
					WORK(0);
				}
#undef WORK
				// process the remaining rows without vectorization
				for (Index i = actual_b_end2; i < actual_b; ++i) {
					if (RK == 4) {
						C0[i] += A0[i] * Bc0[0] + A1[i] * Bc0[1] + A2[i] * Bc0[2] + A3[i] * Bc0[3];
						C1[i] += A0[i] * Bc1[0] + A1[i] * Bc1[1] + A2[i] * Bc1[2] + A3[i] * Bc1[3];
					} else {
						C0[i] += A0[i] * Bc0[0] + A1[i] * Bc0[1];
						C1[i] += A0[i] * Bc1[0] + A1[i] * Bc1[1];
					}
				}

				Bc0 += RK;
				Bc1 += RK;
			} // peeled loop on k
		}	  // peeled loop on the columns j
		// process the last column (we now perform a matrix-vector product)
		if ((n - n_end) > 0) {
			const Scalar* Bc0 = B + (n - 1) * ldb;

			for (Index k = 0; k < d_end; k += RK) {

				// load and expand a 1 x RK block of B
				Packet b00, b10, b20, b30;
				b00 = pset1<Packet>(Bc0[0]);
				b10 = pset1<Packet>(Bc0[1]);
				if (RK == 4)
					b20 = pset1<Packet>(Bc0[2]);
				if (RK == 4)
					b30 = pset1<Packet>(Bc0[3]);

				Packet a0, a1, a2, a3, c0, t0 /*, t1*/;

				const Scalar* A0 = A + ib + (k + 0) * lda;
				const Scalar* A1 = A + ib + (k + 1) * lda;
				const Scalar* A2 = A + ib + (k + 2) * lda;
				const Scalar* A3 = A + ib + (k + 3) * lda;

				Scalar* C0 = C + ib + (n_end)*ldc;

				a0 = pload<Packet>(A0);
				a1 = pload<Packet>(A1);
				if (RK == 4) {
					a2 = pload<Packet>(A2);
					a3 = pload<Packet>(A3);
				} else {
					// workaround "may be used uninitialized in this function" warning
					a2 = a3 = a0;
				}

#define WORK(I)                                                                                                        \
	c0 = pload<Packet>(C0 + i + (I)*PacketSize);                                                                       \
	KMADD(c0, a0, b00, t0)                                                                                             \
	a0 = pload<Packet>(A0 + i + (I + 1) * PacketSize);                                                                 \
	KMADD(c0, a1, b10, t0)                                                                                             \
	a1 = pload<Packet>(A1 + i + (I + 1) * PacketSize);                                                                 \
	if (RK == 4) {                                                                                                     \
		KMADD(c0, a2, b20, t0)                                                                                         \
	}                                                                                                                  \
	if (RK == 4) {                                                                                                     \
		a2 = pload<Packet>(A2 + i + (I + 1) * PacketSize);                                                             \
	}                                                                                                                  \
	if (RK == 4) {                                                                                                     \
		KMADD(c0, a3, b30, t0)                                                                                         \
	}                                                                                                                  \
	if (RK == 4) {                                                                                                     \
		a3 = pload<Packet>(A3 + i + (I + 1) * PacketSize);                                                             \
	}                                                                                                                  \
	pstore(C0 + i + (I)*PacketSize, c0);

				// aggressive vectorization and peeling
				for (Index i = 0; i < actual_b_end1; i += PacketSize * 8) {
					EIGEN_ASM_COMMENT("SPARSELU_GEMML_KERNEL2");
					WORK(0);
					WORK(1);
					WORK(2);
					WORK(3);
					WORK(4);
					WORK(5);
					WORK(6);
					WORK(7);
				}
				// vectorization only
				for (Index i = actual_b_end1; i < actual_b_end2; i += PacketSize) {
					WORK(0);
				}
				// remaining scalars
				for (Index i = actual_b_end2; i < actual_b; ++i) {
					if (RK == 4)
						C0[i] += A0[i] * Bc0[0] + A1[i] * Bc0[1] + A2[i] * Bc0[2] + A3[i] * Bc0[3];
					else
						C0[i] += A0[i] * Bc0[0] + A1[i] * Bc0[1];
				}

				Bc0 += RK;
#undef WORK
			}
		}

		// process the last columns of A, corresponding to the last rows of B
		Index rd = d - d_end;
		if (rd > 0) {
			for (Index j = 0; j < n; ++j) {
				enum
				{
					Alignment = PacketSize > 1 ? Aligned : 0
				};
				typedef Map<Matrix<Scalar, Dynamic, 1>, Alignment> MapVector;
				typedef Map<const Matrix<Scalar, Dynamic, 1>, Alignment> ConstMapVector;
				if (rd == 1)
					MapVector(C + j * ldc + ib, actual_b) +=
						B[0 + d_end + j * ldb] * ConstMapVector(A + (d_end + 0) * lda + ib, actual_b);

				else if (rd == 2)
					MapVector(C + j * ldc + ib, actual_b) +=
						B[0 + d_end + j * ldb] * ConstMapVector(A + (d_end + 0) * lda + ib, actual_b) +
						B[1 + d_end + j * ldb] * ConstMapVector(A + (d_end + 1) * lda + ib, actual_b);

				else
					MapVector(C + j * ldc + ib, actual_b) +=
						B[0 + d_end + j * ldb] * ConstMapVector(A + (d_end + 0) * lda + ib, actual_b) +
						B[1 + d_end + j * ldb] * ConstMapVector(A + (d_end + 1) * lda + ib, actual_b) +
						B[2 + d_end + j * ldb] * ConstMapVector(A + (d_end + 2) * lda + ib, actual_b);
			}
		}

	} // blocking on the rows of A and C
}
#undef KMADD

} // namespace internal

} // namespace Eigen

#endif // EIGEN_SPARSELU_GEMM_KERNEL_H
